package media

import (
	"errors"
	"server/common"
	"sync"
	"time"

	"github.com/jiyeyuran/mediasoup-go/v2"
	"github.com/sirupsen/logrus"
)

type TransportEntry struct {
	TransportId string
	Connected   bool
	Transport   *mediasoup.Transport
	Producers   map[string]*mediasoup.Producer // key: producer_id
	Consumers   map[string]*mediasoup.Consumer // key: consumer_id
}

type RouterEntry struct {
	RouterId   string
	WorkerId   int
	RoomId     string
	Router     *mediasoup.Router
	Transports map[string]TransportEntry // key: transport_id
}

type MediaService struct {
	wsClient  *common.WsClient
	config    Config
	workerMgr *WorkerManager
	serviceId string
	routers   map[string]RouterEntry
	mtx       sync.RWMutex // Mutex to protect access to the routers map
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func NewMediaService() *MediaService {
	return &MediaService{
		wsClient:  nil,
		config:    Config{},
		workerMgr: NewWorkerManager(),
	}
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) Init(config *Config) error {
	if config == nil {
		return errors.New("config cannot be nil")
	}

	var err error
	s.wsClient, err = common.NewWsClient(
		common.ClientConfig{
			URL:            "ws://47.97.17.103:8000",
			Dialer:         nil,
			ReconnectDelay: 5 * time.Second,
			MaxRetries:     5000,
		},
		s.HandleWsMessage,
		s.HandleWsError,
	)
	if err != nil {
		LOG().WithError(err).Error("Failed to create WebSocket client")
		return err
	}

	if err := s.workerMgr.Init(config); err != nil {
		LOG().WithError(err).Error("Failed to initialize worker manager")
		return err
	}

	s.serviceId = common.GenerateRandomStr(16)

	s.routers = make(map[string]RouterEntry)

	s.config = *config

	return nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) Start() {
	if err := s.wsClient.Start(); err != nil {
		LOG().WithError(err).Error("Failed to connect WebSocket client")
		return
	}

	go func() {
		ticker := time.NewTicker(60 * time.Second)
		defer ticker.Stop()
		for {
			registerReq := common.Message{
				Cmd: common.REGISTER_SERVICE_REQ,
				Rid: common.GenerateRandomStr(8),
				Data: common.RegisterServiceReqData{
					ServiceId:   s.serviceId,
					ServiceType: common.SERVICE_TYPE_MEDIA,
					Workers:     s.workerMgr.GetWorkers(),
				},
			}
			if err := s.wsClient.Send(registerReq); err != nil {
				LOG().WithError(err).Error("Failed to send WebSocket message")
			} else {
				LOG().Info("Register request sent")
			}
			<-ticker.C
		}
	}()

	LOG().Info("Media service started and WebSocket client connected")

	select {} // 阻塞调用线程
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) HandleWsMessage(msg common.Message) {
	LOG().Infof("Handling WebSocket message: %s", msg)
	switch msg.Cmd {
	case common.REGISTER_SERVICE_RES:
		s.handleRegisterServiceRes(msg)
	case common.CREATE_ROUTER_REQ:
		s.handleCreateRouterReq(msg)
	case common.CLOSE_ROUTER_REQ:
		s.handleCloseRouterReq(msg)
	case common.CREATE_TRANSPORT_REQ:
		s.handleCreateTransportReq(msg)
	case common.CLOSE_TRANSPORT_REQ:
		s.handleCloseTransportReq(msg)
	case common.CONNECT_TRANSPORT_REQ:
		s.handleConnectTransportReq(msg)
	case common.CREATE_PRODUCER_REQ:
		s.handleCreateProducerReq(msg)
	case common.CLOSE_PRODUCER_REQ:
		s.handleCloseProducerReq(msg)
	case common.CREATE_CONSUMER_REQ:
		s.handleCreateConsumerReq(msg)
	case common.CLOSE_CONSUMER_REQ:
		s.handleCloseConsumerReq(msg)
	default:
		LOG().Warnf("Unhandled WebSocket message command: %s", msg.Cmd)
	}
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) HandleWsError(err error) {
	LOG().Errorf("WebSocket error: %v", err)
	// 这里可以添加错误处理逻辑
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CreateRouter(rid string, reqData common.CreateRouterReqData) (
	*mediasoup.Router, error) {
	s.mtx.Lock()
	defer s.mtx.Unlock()

	worker, err := s.workerMgr.GetWorker(reqData.WorkerId)
	if err != nil {
		LOG(rid).WithError(err).Error("failed to get worker")
		return nil, err
	}

	router, err := worker.CreateRouter(s.config.Mediasoup.RouterOptions)
	if err != nil {
		LOG(rid).WithError(err).Error("failed to create router")
		return nil, err
	}

	if _, exists := s.routers[router.Id()]; exists {
		LOG(rid).Error("router already exists", "routerId", router.Id())
		return nil, errors.New("router already exists")
	}

	s.routers[router.Id()] = RouterEntry{
		RouterId:   router.Id(),
		WorkerId:   worker.Pid(),
		RoomId:     reqData.RoomId,
		Router:     router,
		Transports: make(map[string]TransportEntry),
	}

	LOG(rid).WithFields(logrus.Fields{
		"id":   router.Id(),
		"size": len(s.routers),
	}).Info("Router created successfully")

	return router, nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CloseRouter(rid string, routerId string, workerId int) error {
	s.mtx.Lock()
	defer s.mtx.Unlock()

	routerEntry, exists := s.routers[routerId]
	if !exists {
		err := errors.New("router not found")
		LOG(rid).WithField("routerId", routerId).WithError(err).Error("router does not exist")
		return err
	}

	if routerEntry.WorkerId != workerId {
		err := errors.New("worker ID mismatch")
		LOG(rid).WithField("workerId", workerId).WithError(err).Error("worker ID does not match")
		return err
	}

	if err := routerEntry.Router.Close(); err != nil {
		LOG(rid).WithError(err).Error("failed to close router")
		return err
	}

	delete(s.routers, routerId)

	LOG(rid).WithField("routerId", routerId).Info("Router closed successfully")

	return nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CreateTransport(rid string, reqData common.CreateTransportReqData) (
	*mediasoup.Transport, error) {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	routerEntry, exists := s.routers[reqData.RouterId]
	if !exists {
		LOG(rid).Error("Router not found", "routerId", reqData.RouterId)
		return nil, errors.New("router not found")
	}

	webRtcTransportOptions := &mediasoup.WebRtcTransportOptions{
		AppData: mediasoup.H{
			"producing": reqData.Producing,
			"consuming": reqData.Consuming,
		},
	}
	common.Clone(webRtcTransportOptions, s.config.Mediasoup.WebRtcTransportOptions)

	webRtcTransportOptions.EnableSctp = true

	if reqData.SctpCaps != nil {
		webRtcTransportOptions.NumSctpStreams = &reqData.SctpCaps.NumStreams
	}

	if reqData.ForceTcp {
		webRtcTransportOptions.EnableUdp = common.Ref(false)
		webRtcTransportOptions.EnableTcp = true
	}

	transport, err := routerEntry.Router.CreateWebRtcTransport(webRtcTransportOptions)
	if err != nil {
		LOG(rid).WithError(err).Error("Failed to create WebRTC transport")
		return nil, err
	}

	routerEntry.Transports[transport.Id()] = TransportEntry{
		TransportId: transport.Id(),
		Transport:   transport,
		Connected:   false,
		Producers:   make(map[string]*mediasoup.Producer),
		Consumers:   make(map[string]*mediasoup.Consumer),
	}

	LOG(rid).WithFields(logrus.Fields{
		"routerId":    reqData.RouterId,
		"transportId": transport.Id(),
		"size":        len(routerEntry.Transports),
	}).Info("WebRTC transport created successfully")

	return transport, nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) ConnectTransport(rid string, reqData common.ConnectTransportReqData) error {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	transportEntry, err := s.GetTransport(reqData.RouterId, reqData.TransportId)
	if err != nil {
		LOG(rid).WithError(err).Error("Failed to get transport")
		return err
	}

	if transportEntry.Connected {
		LOG(rid).Warn("Transport is already connected")
		return nil
	}

	if err := transportEntry.Transport.Connect(&mediasoup.TransportConnectOptions{
		DtlsParameters: reqData.DtlsParameters,
	}); err != nil {
		LOG(rid).WithError(err).Error("Failed to connect transport")
		return err
	}

	transportEntry.Connected = true
	LOG(rid).WithField("id", transportEntry.TransportId).Info("Transport connected successfully")

	return nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CloseTransport(routerId, transportId string) error {
	s.mtx.Lock()
	defer s.mtx.Unlock()

	routerEntry, exists := s.routers[routerId]
	if !exists {
		return errors.New("router not found")
	}

	transportEntry, exists := routerEntry.Transports[transportId]
	if !exists {
		return errors.New("transport not found")
	}

	if err := transportEntry.Transport.Close(); err != nil {
		LOG().WithError(err).Error("Failed to close transport")
		return err
	}

	delete(routerEntry.Transports, transportId)
	LOG().WithField("id", transportId).Info("Transport closed successfully")

	return nil
}

// -------------------------------------------------------------------------------------------------
// Note: Must lock the mutex before calling this function to ensure thread safety.
// -------------------------------------------------------------------------------------------------
func (s *MediaService) GetTransport(routerId, transportId string) (*TransportEntry, error) {
	routerEntry, exists := s.routers[routerId]
	if !exists {
		return nil, errors.New("router not found")
	}

	transport, exists := routerEntry.Transports[transportId]
	if !exists {
		return nil, errors.New("transport not found")
	}

	return &transport, nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) RouterExists(routerId string) bool {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	_, exists := s.routers[routerId]
	return exists
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) TransportExists(routerId, transportId string) bool {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	routerEntry, exists := s.routers[routerId]
	if !exists {
		return false
	}

	_, exists = routerEntry.Transports[transportId]
	return exists
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CreateProducer(rid string, reqData common.CreateProducerReqData) (
	*mediasoup.Producer, error) {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	transport, err := s.GetTransport(reqData.RouterId, reqData.TransportId)
	if err != nil {
		LOG(rid).WithError(err).Error("failed to get transport")
		return nil, err
	}

	producer, err := transport.Transport.Produce(&mediasoup.ProducerOptions{
		Kind:          reqData.Kind,
		RtpParameters: reqData.RtpParameters,
		AppData:       reqData.AppData,
	})
	if err != nil {
		LOG(rid).WithError(err).Error("failed to produce media")
		return nil, err
	}

	transport.Producers[producer.Id()] = producer

	LOG(rid).WithFields(logrus.Fields{
		"routerId":      reqData.RouterId,
		"transportId":   reqData.TransportId,
		"producerId":    producer.Id(),
		"size":          len(transport.Producers),
		"rtpParameters": common.ToJson(producer.RtpParameters()),
	}).Info("Producer created successfully")

	return producer, nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CloseProducer(rid string, reqData common.CloseProducerReqData) error {
	s.mtx.Lock()
	defer s.mtx.Unlock()

	transport, err := s.GetTransport(reqData.RouterId, reqData.TransportId)
	if err != nil {
		LOG(rid).WithError(err).Error("failed to get transport")
		return err
	}

	producer, exists := transport.Producers[reqData.ProducerId]
	if !exists {
		err := errors.New("producer not found")
		LOG(rid).WithField("producerId", reqData.ProducerId).WithError(err).Error("producer does not exist")
		return err
	}

	if err := producer.Close(); err != nil {
		LOG(rid).WithError(err).Error("failed to close producer")
		return err
	}

	delete(transport.Producers, reqData.ProducerId)

	LOG(rid).WithField("producerId", reqData.ProducerId).Info("Producer closed successfully")

	return nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) ProducerExists(routerId, producerId string) bool {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	routerEntry, exists := s.routers[routerId]
	if !exists {
		return false
	}

	for _, transport := range routerEntry.Transports {
		if _, exists := transport.Producers[producerId]; exists {
			return true
		}
	}

	return false
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CreateConsumer(rid string, reqData common.CreateConsumerReqData) (
	*mediasoup.Consumer, error) {
	s.mtx.RLock()
	transport, err := s.GetTransport(reqData.RouterId, reqData.TransportId)
	if err != nil {
		LOG(rid).WithFields(logrus.Fields{
			"routerId":    reqData.RouterId,
			"transportId": reqData.TransportId,
		}).WithError(err).Error("failed to get transport")
		return nil, err
	}
	s.mtx.RUnlock()

	if !s.ProducerExists(reqData.RouterId, reqData.ProducerId) {
		err := errors.New("producer not found")
		LOG(rid).WithField("producerId", reqData.ProducerId).WithError(err).Error("producer does not exist")
		return nil, err
	}

	consumer, err := transport.Transport.Consume(&mediasoup.ConsumerOptions{
		ProducerId:      reqData.ProducerId,
		RtpCapabilities: reqData.RtpCaps,
		AppData:         reqData.AppData,
		Paused:          false,
		EnableRtx:       common.Ref(true),
		IgnoreDtx:       true,
	})
	if err != nil {
		LOG(rid).WithError(err).Error("failed to consume media")
		return nil, err
	}

	{
		s.mtx.RLock()
		transport.Consumers[consumer.Id()] = consumer
		s.mtx.RUnlock()
	}

	LOG(rid).WithFields(logrus.Fields{
		"routerId":      reqData.RouterId,
		"transportId":   reqData.TransportId,
		"producerId":    reqData.ProducerId,
		"consumerId":    consumer.Id(),
		"rtpParameters": common.ToJson(reqData.RtpCaps),
		"appData":       common.ToJson(reqData.AppData),
		"size":          len(transport.Consumers),
	}).Info("Consumer created successfully")

	return consumer, nil
}

// -------------------------------------------------------------------------------------------------
//
// -------------------------------------------------------------------------------------------------
func (s *MediaService) CloseConsumer(rid string, reqData common.CloseConsumerReqData) error {
	s.mtx.Lock()
	defer s.mtx.Unlock()

	transport, err := s.GetTransport(reqData.RouterId, reqData.TransportId)
	if err != nil {
		LOG(rid).WithError(err).Error("failed to get transport")
		return err
	}

	consumer, exists := transport.Consumers[reqData.ConsumerId]
	if !exists {
		err := errors.New("consumer not found")
		LOG(rid).WithField("consumerId", reqData.ConsumerId).WithError(err).Error("consumer does not exist")
		return err
	}

	if err := consumer.Close(); err != nil {
		LOG(rid).WithError(err).Error("failed to close consumer")
		return err
	}

	delete(transport.Consumers, reqData.ConsumerId)

	LOG(rid).WithField("consumerId", reqData.ConsumerId).Info("Consumer closed successfully")

	return nil
}
